import os
import pandas as pd
import statsmodels.api as sm
import matplotlib.pyplot as plt

join = os.path.join
dirname = os.path.dirname
CURRENT_PATH = dirname(os.path.realpath(__file__))
DATA_PATH = join(dirname(CURRENT_PATH) , 'DataAssets')
SAVE_PATH = ''

'''
遍历所有序列，输出可视化输出周期性分解的结果
发现大部分序列的残差项过大，分解效果不佳，后续不使用seasonal因子

input :pd.series 时序数据
output : csv文件 周期性分解后信息
output : jpg文件 周期性分解后的图片
'''

class Seasonal_decompose_analysis:
    def __init__(self, endog, ts_name, seasonal_period ,model_type, if_save = False  , info_SAVE_PATH = None , two_sided=False):

        self.data_ts = endog
        self.two_sided = two_sided
        self.ts_name = ts_name
        self.model_type = model_type # 'additive' 'multiplicate'
        

        if if_save:
            self.info_SAVE_PATH = info_SAVE_PATH
        else:
            self.info_SAVE_PATH = None

        # 原始ts拆解后的三部分
        self.decom_ts = None
        
        self.seasonal_period = seasonal_period
        
        
        if self.data_ts.shape[0] < 6 and self.seasonal_period >= 3:
            raise('self.data_ts.shape[0] < 6 but self.seasonal_period >= 3')
        elif self.data_ts.shape[0] < 12 and self.seasonal_period >= 6:
            raise('self.data_ts.shape[0] < 12 but self.seasonal_period >= 6')
        elif self.data_ts.shape[0] < 24 and self.seasonal_period == 12:
            raise('self.data_ts.shape[0] < 24 but self.seasonal_period >= 12')
        elif self.seasonal_period > 12:
            raise('self.seasonal_period > 12')


    def decompose_ts(self):
        """
        使用季节性拆解方法对原序列进行分拆
        :return: 拆解后的三个部分
        """
        decomposition = sm.tsa.seasonal_decompose(self.data_ts, model=self.model_type, period=self.seasonal_period, \
                                                  two_sided=self.two_sided)

        trend_ts = decomposition.trend.dropna()
        seasonal_ts = decomposition.seasonal.dropna()
        residual_ts = decomposition.resid.dropna()

        self.decom_ts = {'trend': trend_ts,
                         'seasonal': seasonal_ts,
                         'residual': residual_ts}

        
        if self.info_SAVE_PATH:
            decom_ts_df = pd.DataFrame(self.decom_ts, index=[0])
            decom_ts_df_describe = decom_ts_df.describe()
            
            with pd.ExcelWriter(join(self.info_SAVE_PATH , self.ts_name + '_decompose.xlsx')) as writer:
                decom_ts_df.to_excel(writer, sheet_name='decompose_detail')
                decom_ts_df_describe.to_excel(writer, sheet_name='decompose_summary')

    def decompose_draw(self ):
        """
        对分解后的序列进行绘图
        :return:
        """
        if self.decom_ts is None:
            self.decompose_ts()

        _, _ = plt.subplots(figsize=(12, 8))
        plt.subplot(411)
        plt.plot(self.data_ts, label='Original')
        plt.legend(loc='best')

        plt.subplot(412)
        plt.plot(self.decom_ts['trend'], label='Trend')
        plt.legend(loc='best')

        plt.subplot(413)
        plt.plot(self.decom_ts['seasonal'], label='Seasonality')
        plt.legend(loc='best')

        plt.subplot(414)
        plt.plot(self.decom_ts['residual'], label='Residuals')
        plt.legend(loc='best')

        plt.tight_layout()
        if self.info_SAVE_PATH:
            plt.savefig(join(self.info_SAVE_PATH , self.ts_name + '.jpg'))
        plt.show()
        plt.close()


if __name__ == '__main__':
    cur_name = 'test_arima'  # 自变量

    cur_file = pd.read_excel(join(DATA_PATH, ''))

    cur_file_one = cur_file.loc[cur_file['asc_code'] == 0000000, 'value']

    out_of_var = pd.Series([0, 1, 0])
    arima_model_one = Seasonal_decompose_analysis(endog=cur_file_one,\
                                         ts_name='model_test',\
                                         seasonal_period=4,
                                         model_type= 'additive',
                                         if_save = True,\
                                         info_SAVE_PATH = SAVE_PATH)

    arima_model_one.decompose_draw()
